/*
 * Copyright 2021 Rockchip Electronics Co. LTD
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc.hpp>

#include "rtspSend.hpp"
#ifdef __cplusplus
#if __cplusplus
extern "C" {
#endif
#endif /* End of #ifdef __cplusplus */

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <pthread.h>
#include <pthread.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/prctl.h>
#include <time.h>
#include <unistd.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <linux/serial.h>
#include <linux/videodev2.h>
#include <sys/mman.h>
//#include "rtsp_demo.h"
#include "sample_comm.h"

#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/opt.h>
#include <libavutil/imgutils.h>
#ifdef __cplusplus
#if __cplusplus
}
#endif
#endif /* End of #ifdef __cplusplus */

typedef struct _rkMpiCtx {
	SAMPLE_VI_CTX_S vi;
	SAMPLE_VO_CTX_S vo;
	SAMPLE_VPSS_CTX_S vpss;
	SAMPLE_VENC_CTX_S venc;
	SAMPLE_RGN_CTX_S rgn[2];
} SAMPLE_MPI_CTX_S;

static SAMPLE_MPI_CTX_S *g_ctx;
static bool quit = false;
int video_width = 1920;
int video_height = 1080;
// static COMPRESS_MODE_E g_compressMode = COMPRESS_AFBC_16x16;
static COMPRESS_MODE_E g_compressMode = COMPRESS_MODE_NONE;
MB_POOL pool = MB_INVALID_POOLID;
float fps_use = 30;

void signalHandler(int sig) {
	printf("Signal usbgetimg, exiting...\n");
	pthread_exit(NULL);
}
RK_S32 s32BitRate = 2*1024*1024;
AVFormatContext* outputFormatContext = nullptr;
AVStream* stream = nullptr;
bool isInitialized = false;
bool isPushing= false;
int64_t pts = 0;
int64_t dts = 0;
int frameCount = 0;
AVRational timeBase = {1, 30};
char ffmpegUrl[200] = {0};
bool codecis;
int retryCount = 0;
int maxRetries = 5;
char errbuf[AV_ERROR_MAX_STRING_SIZE];
 // 开始推流
bool start() {
	if (!isInitialized) {
		std::cerr << "推流器未初始化" << std::endl;
		return false;
	}
	
	if (isPushing) {
		std::cerr << "推流器已在运行" << std::endl;
		return false;
	}
	
	isPushing = true;
	return true;
}

// 停止推流
void stop() {
	isPushing = false;
}
 // 检查推流器是否正在运行
bool isRunning() {
	return isPushing;
}

// 关闭流连接
void closeStream() {
	// if (!isInitialized) {
	// 	return;
	// }

	// 写入文件尾
	if (outputFormatContext) {
		av_write_trailer(outputFormatContext);
	}

	// 清理资源
	if (outputFormatContext && !(outputFormatContext->oformat->flags & AVFMT_NOFILE)) {
		avio_closep(&outputFormatContext->pb);
	}
	if (outputFormatContext) {
		avformat_free_context(outputFormatContext);
	}
	
	outputFormatContext = nullptr;
	stream = nullptr;
	isInitialized = false;
}
// 打开流连接
bool openStream() {
	closeStream();
	int ret = 0;
	//printf("程序链接的 FFmpeg 版本: %d\n", avformat_version());
	// 注册所有FFmpeg组件
	//av_register_all();
	//avcodec_register_all();
	ret = avformat_network_init();
	if (ret < 0) {
        char errbuf[AV_ERROR_MAX_STRING_SIZE];
        av_strerror(ret, errbuf, AV_ERROR_MAX_STRING_SIZE);
        fprintf(stderr, "网络初始化失败: %s\n", errbuf);
        return false;
    }
#if 0
	const AVOutputFormat *ofmt = NULL;
    void *iter = NULL;
	printf("FFmpeg支持的输出格式：\n");
    while ((ofmt = av_muxer_iterate(&iter))) {
        printf("名称: %-15s 长名称: %s\n", ofmt->name, ofmt->long_name);
        printf("MIME类型: %s\n", ofmt->mime_type ? ofmt->mime_type : "N/A");
        printf("扩展名: %s\n", ofmt->extensions ? ofmt->extensions : "N/A");
        printf("----------------------------------------\n");
    }
#endif
	// 查找RTSP muxer
	// 创建输出格式上下文
	ret = avformat_alloc_output_context2(&outputFormatContext, nullptr, "rtsp", ffmpegUrl);
	if (ret < 0) {
		char errbuf[AV_ERROR_MAX_STRING_SIZE];
        av_strerror(ret, errbuf, AV_ERROR_MAX_STRING_SIZE);
		//std::cerr << "无法创建输出格式上下文: " << av_err2str(ret) << std::endl;
		std::cerr << "无法创建输出格式上下文: " << errbuf << std::endl;
		return false;
	}

	// 根据输入编码设置输出流参数
	stream = avformat_new_stream(outputFormatContext, nullptr);
	if (!stream) {
		std::cerr << "无法创建新流" << std::endl;
		return false;
	}

	// 设置编码器ID和参数
	if (!codecis) {
		stream->codecpar->codec_id = AV_CODEC_ID_H264;
	} else if (codecis) {
		stream->codecpar->codec_id = AV_CODEC_ID_HEVC;
	} 
	timeBase = {1, (int)fps_use};
	stream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
	stream->codecpar->format = AV_PIX_FMT_YUV420P;
	stream->time_base = timeBase;
	stream->avg_frame_rate = {(int)fps_use, 1};
	
	// 设置视频尺寸（关键修改点）
	stream->codecpar->width = video_width;
	stream->codecpar->height = video_height;
	stream->codecpar->format = AV_PIX_FMT_YUV420P;
	
	// 设置合理的比特率估计值，帮助rtsp-simple-server优化流传输
	if (!codecis) {
		stream->codecpar->bit_rate = s32BitRate * 1024; // 2Mbps
	} else { // H.265
		stream->codecpar->bit_rate = s32BitRate * 1024; // 1.5Mbps (通常H.265效率更高)
	}

	// 设置RTSP特定选项
	AVDictionary* options = nullptr;
	// 使用TCP传输，更可靠
	av_dict_set(&options, "rtsp_transport", "tcp", 0);
	// 设置合理的缓冲区大小
	av_dict_set(&options, "buffer_size", "4096000", 0);
	// 设置超时时间
	av_dict_set(&options, "timeout", "5000000", 0); // 5秒

	// 打开输出URL
	if (!(outputFormatContext->oformat->flags & AVFMT_NOFILE)) {
		ret = avio_open2(&outputFormatContext->pb, ffmpegUrl, AVIO_FLAG_WRITE, nullptr, &options);
		if (ret < 0) {
			std::cerr << "无法打开输出URL: " << ret << std::endl;
			av_dict_free(&options);
			return false;
		}
	}
	av_dict_free(&options);

	// 写入文件头
	ret = avformat_write_header(outputFormatContext, nullptr);
	if (ret < 0) {
		std::cerr << "无法写入文件头: " << ret << std::endl;
		return false;
	}

	isInitialized = true;
	std::cout << "RTSP推流器初始化成功，目标URL: " << ffmpegUrl << std::endl;
	return true;
}

/// 推送一帧H.264/H.265裸码流数据
bool ffmpeg_pushFrame(const uint8_t* data, size_t size, bool isKeyFrame = false) {
	if (!isInitialized || !isPushing || !data || size == 0) {
		return false;
	}
	// 准备AVPacket
	AVPacket* packet = av_packet_alloc();
	if (!packet) {
		std::cerr << "无法分配数据包" << std::endl;
		return false;
	}

	// 复制数据到packet
	int ret = av_new_packet(packet, size);
	if (ret < 0) {
		std::cerr << "无法分配数据包内存: " << ret << std::endl;
		av_packet_free(&packet);
		return false;
	}
	memcpy(packet->data, data, size);

	// 设置时间戳和流索引
	packet->stream_index = stream->index;
	packet->pts = pts;
	packet->dts = dts;
	packet->duration = av_rescale_q(1, timeBase, stream->time_base);
	// 设置关键帧标志
	if (isKeyFrame) {
		packet->flags |= AV_PKT_FLAG_KEY;
	} else {
		// 如果未指定是否为关键帧，尝试通过NAL单元类型检测
		if (size >= 5) {
			// H.264 NAL单元类型检测
			if (stream->codecpar->codec_id == AV_CODEC_ID_H264) {
				uint8_t nal_type = data[4] & 0x1F;
				if (nal_type == 7 || nal_type == 8 || nal_type == 5) { // SPS, PPS或IDR
					packet->flags |= AV_PKT_FLAG_KEY;
				}
			}
			// H.265 NAL单元类型检测
			else if (stream->codecpar->codec_id == AV_CODEC_ID_HEVC) {
				uint8_t nal_type = (data[4] >> 1) & 0x3F;
				if (nal_type == 32 || nal_type == 33 || nal_type == 19 || nal_type == 20) { // VPS, SPS或IDR
					packet->flags |= AV_PKT_FLAG_KEY;
				}
			}
		}
	}

	// 发送数据包
	ret = av_interleaved_write_frame(outputFormatContext, packet);
	if (ret < 0) {
		std::cerr << "写入数据包失败: " << ret << std::endl;
		
		// 可能是连接断开，标记需要重新连接
		closeStream();
		
		av_packet_free(&packet);
		return false;
	}
	// 更新时间戳
	pts += packet->duration;
	dts += packet->duration;
	frameCount++;

	// 释放数据包
	av_packet_free(&packet);

	return true;
}

// 关闭推流器
void ffmpeg_close() {
	if (!isInitialized) {
		return;
	}

	// 写入文件尾
	av_write_trailer(outputFormatContext);

	// 清理资源
	if (outputFormatContext && !(outputFormatContext->oformat->flags & AVFMT_NOFILE)) {
		avio_closep(&outputFormatContext->pb);
	}
	avformat_free_context(outputFormatContext);
	avformat_network_deinit();

	outputFormatContext = nullptr;
	isInitialized = false;

	std::cout << "RTSP推流器已关闭，共发送 " << frameCount << " 帧" << std::endl;
}

/******************************************************************************
* function : venc thread
******************************************************************************/
static void *venc_get_stream(void *pArgs) {
	SAMPLE_VENC_CTX_S *ctx = (SAMPLE_VENC_CTX_S *)(pArgs);
	RK_S32 s32Ret = RK_FAILURE;
	void *pData = RK_NULL;
	RK_S32 loopCount = 0;
	signal(SIGUSR2, signalHandler); // 设置信号处理函数
	//printf("in venc_get_stream\n");
	// sleep(1);
	while (!quit) {
		//printf("in venc_get_stream while\n");
		s32Ret = SAMPLE_COMM_VENC_GetStream(ctx, &pData);
		if (s32Ret == RK_SUCCESS) {
			// exit when complete
			if (ctx->s32loopCount > 0) {
				if (loopCount >= ctx->s32loopCount) {
					SAMPLE_COMM_VENC_ReleaseStream(ctx);
					quit = true;
					break;
				}
			}
			
			if (quit)
			{
				SAMPLE_COMM_VENC_ReleaseStream(ctx);
				break;
			}
			ffmpeg_pushFrame((uint8_t*)pData, ctx->stFrame.pstPack->u32Len);
			
			// printf("chn:%d, loopCount:%d wd:%d ret %d\n", ctx->s32ChnId, loopCount,ctx->stFrame.pstPack->u32Len, ret);

			SAMPLE_COMM_VENC_ReleaseStream(ctx);
			loopCount++;
		}
		usleep(1000);
	}
	return RK_NULL;
}

YourStreamer::YourStreamer()
{
}
YourStreamer::~YourStreamer()
{
}
//!用于初始化视频推流，仅调用一次
/*!
	\param name:推流的子地址名称，网络地址为:rtsp://YourIp:port/name
	\param size:保存分辨率，如cv::Size(1920，1080)
	\param use_h265:是否使用h265编码
	\param port:端口号，默认是8554
	\param bitrate:推流时候的码率，默认为2兆，单位是兆
	返回是否初始化成功
*/
bool YourStreamer::setup(std::string name, cv::Size size, bool use_h265, std::string port, int bitrate, float fps)
{
	//获取参数
    this->name =name;
    this->size =size;
    this->use_h265 =use_h265;
    this->port = port;
    this->bitrate =bitrate;
	fps_use = fps;
	int ret = 0;
	int profile = 0;
	video_width = size.width;
	video_height = size.height;
	s32BitRate = bitrate * 1024;
	RK_S32 s32loopCnt = -1;
	SAMPLE_MPI_CTX_S* ctx;
	CODEC_TYPE_E enCodecType = RK_CODEC_TYPE_H264;
	VENC_RC_MODE_E enRcMode = VENC_RC_MODE_H264CBR;
	codecis = use_h265;
	//init rockit
	ctx = (SAMPLE_MPI_CTX_S*)(malloc(sizeof(SAMPLE_MPI_CTX_S)));
	memset(ctx, 0, sizeof(SAMPLE_MPI_CTX_S));
	g_ctx = ctx;
	if (RK_MPI_SYS_Init() != RK_SUCCESS) {
		RK_MPI_SYS_Exit();
		if (ctx) {
			free(ctx);
			ctx = RK_NULL;
		}
		return false;
	}

	//init mem
	
	MB_POOL_CONFIG_S stMbPoolCfg;
	memset(&stMbPoolCfg, 0, sizeof(MB_POOL_CONFIG_S));
	stMbPoolCfg.u64MBSize = video_width * video_height * 2;
	stMbPoolCfg.u32MBCnt = 60;
	stMbPoolCfg.enAllocType = MB_ALLOC_TYPE_DMA;
	pool = RK_MPI_MB_CreatePool(&stMbPoolCfg);
	if (pool == MB_INVALID_POOLID)
	{
		printf("mem init error\n");
		ctx->venc.pool = pool;
		return false;
	}
	ctx->venc.pool = pool;

	// Init VENC[0]
	if (use_h265)
	{
		enCodecType = RK_CODEC_TYPE_H265;
		enRcMode = VENC_RC_MODE_H264CBR;
		profile = 0;
	}
	else
	{
		enCodecType = RK_CODEC_TYPE_H264;
		enRcMode = VENC_RC_MODE_H264CBR;
		profile = 100;
	}
	ctx->venc.s32ChnId = 0;
	ctx->venc.u32Width = video_width;
	ctx->venc.u32Height = video_height;
	ctx->venc.u32Fps = fps_use;
	ctx->venc.u32Gop = 60;
	ctx->venc.u32BitRate = s32BitRate;
	ctx->venc.enCodecType = enCodecType;
	ctx->venc.enRcMode = enRcMode;
	ctx->venc.getStreamCbFunc = venc_get_stream;
	ctx->venc.s32loopCount = s32loopCnt;
	//ctx->venc.dstFilePath = "test";//pOutPathVenc; guo
	// H264  66：Baseline  77：Main Profile 100：High Profile
	// H265  0：Main Profile  1：Main 10 Profile
	// MJPEG 0：Baseline
	ctx->venc.stChnAttr.stVencAttr.u32Profile = profile;
	ctx->venc.stChnAttr.stGopAttr.enGopMode = VENC_GOPMODE_NORMALP; // VENC_GOPMODE_SMARTP
	ret = SAMPLE_COMM_VENC_CreateChn(&ctx->venc);
	if (ret != RK_SUCCESS)
	{
		printf("init h264enc error\n");
		return false;
	}

	snprintf(ffmpegUrl, sizeof(ffmpegUrl), "rtsp://127.0.0.1:%s%s", port.c_str(), name.c_str());
	//init rtsp
	// ret = openStream();
	// if(!ret)
	// {
	// 	printf("openStream error \n");
	// 	return false;
	// }
	// ret = start();
	// if(!ret)
	// {
	// 	printf("openStream error \n");
	// 	return false;
	// }
	
	
	return true;
}

//!判断推流通道是否已经打开
bool YourStreamer::isOpened()
{
	if(!isInitialized)
	{
		//init rtsp
		bool ret = openStream();
		if(!ret)
		{
			printf("openStream error \n");
			return false;
		}
		ret = start();
		if(!ret)
		{
			printf("openStream error \n");
			return false;
		}
	}
	return true;
}
struct timeval start_time, end_time;
float timeuse = 0.0;
//!用于推流每一帧图像，在循环中被调用
/*!
	\param image:待推流的视频帧
*/
void YourStreamer::stream(cv::Mat image)
{
	cv::Mat getimg = image.clone();
	//默认image格式为BGR
	//输入图像与rtsp初始化时设置的分辨率不一致，guo:确认是否需要缩放为一致
	if (image.cols != video_width || image.rows != video_height)
	{
		resize(getimg, getimg, Size(video_width, video_height));
		// printf("rtsp input img %dx%d  init size %dx%d\n", image.cols, image.rows, video_width, video_height);
		//return;
	}
	cv::cvtColor(getimg, getimg, cv::COLOR_BGR2YUV_I420);//cv::COLOR_BGR2GRAY);//
	
	
	SAMPLE_COMM_VENC_SendStream(&g_ctx->venc, getimg.data, video_width, video_height,video_width * video_height * 3 / 2, g_compressMode);

	return;
}

//!释放资源
void YourStreamer::release()
{
	int ret = 0;
	quit = true;
	//关闭线程
	//pthread_kill(g_ctx->venc.getStreamThread, SIGUSR2);
	printf("kill enc thread\n");
	if (g_ctx->venc.getStreamCbFunc)
	{
		pthread_join(g_ctx->venc.getStreamThread, NULL); // 等待线程退出
	}
	
	//关闭解码器
	printf("stop enc\n");
	ret = SAMPLE_COMM_VENC_DestroyChn(&g_ctx->venc);
	if (ret != RK_SUCCESS)
	{
		printf("stop enc error!\n");
	}
	//关闭缓存
	printf("stop mem\n");
	if (pool != MB_INVALID_POOLID)
	{
		RK_MPI_MB_DestroyPool(pool);
		pool = MB_INVALID_POOLID;
	}
	//关闭rkmpi
	printf("stop rkmpi\n");
	if (g_ctx) {
		free(g_ctx);
		g_ctx = RK_NULL;
		RK_MPI_SYS_Exit();
	}
	
	// 清理资源
	std::cout << "Stopping pipeline..." << std::endl;
	//if (source_id) g_source_remove(source_id);
	ffmpeg_close();
	printf("to stop end\n");
	return;
}


